"""
python fujian2_CubicSplineFix_gragh.py
"""
import os
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.interpolate import CubicSpline
from datetime import datetime, timedelta

# 设置相对路径
input_path = "../fujian/fujian2/groupByCategory"
output_path = "../fujian/fujian2/CubicSplineFix"
output_json_path = os.path.join(output_path, "categoryJSON")
output_graph_path = os.path.join(output_path, "graph")

os.makedirs(output_json_path, exist_ok=True)
os.makedirs(output_graph_path, exist_ok=True)

# 遍历输入目录中的所有 JSON 文件
json_files = [f for f in os.listdir(input_path) if f.endswith('.json')]

for json_file in json_files:
    file_path = os.path.join(input_path, json_file)
    
    # 读取 JSON 数据
    with open(file_path, 'r', encoding='utf-8') as f:
        sales_data = json.load(f)
    
    # 将数据转换为 DataFrame 并进行日期处理
    df = pd.DataFrame(sales_data)
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)

    # 创建完整日期范围（2022/7/1 - 2023/6/30）
    full_range = pd.date_range(start='2022-07-01', end='2023-06-30')
    full_df = pd.DataFrame(index=full_range)

    # 合并数据并进行三次样条插值
    merged_df = full_df.join(df, how='left')
    # merged_df['sales_interp'] = merged_df['sales'].interpolate(method='cubic')
    merged_df['sales_interp'] = merged_df['sales'].interpolate(method='spline', order=3)


    # 保存补全后的数据到新的 JSON 文件
    filled_data = [{'date': (date + timedelta(days=1)).strftime('%Y/%m/%d'), 'sales': round(sales, 2)}
                   for date, sales in zip(merged_df.index, merged_df['sales_interp'])]
    new_json_file = os.path.join(output_json_path, json_file)
    with open(new_json_file, 'w', encoding='utf-8') as f:
        json.dump(filled_data, f, ensure_ascii=False, indent=4)

    # 绘制插值图
    x_original = (merged_df.index - merged_df.index[0]).days + 1  # 计算从day1开始的天数
    y_original = merged_df['sales'].to_numpy()
    y_interp = merged_df['sales_interp'].to_numpy()
    cs = CubicSpline(x_original[~np.isnan(y_original)], y_original[~np.isnan(y_original)])  # 创建插值模型

    # 创建插值图形
    plt.figure(figsize=(14, 7))
    plt.scatter(x_original, y_original, color='red', label='Original Data', zorder=5)
    plt.scatter(x_original, y_interp, color='blue', label='Interpolated Data', zorder=5)
    
    # 插值曲线
    x_dense = np.linspace(x_original[0], x_original[-1], 1000)
    plt.plot(x_dense, cs(x_dense), color='green', label='Cubic Spline Interpolation', linewidth=2)

    # 图形设置
    plt.title(f'Sales Data with Cubic Spline Interpolation for {json_file}')
    plt.xlabel('Days from 2022-07-01')
    plt.ylabel('Sales')
    plt.legend()
    plt.grid()

    # 保存图像
    graph_file = os.path.join(output_graph_path, f"graph_{json_file.split('.')[0]}.png")
    plt.savefig(graph_file)
    plt.close()

print("插值和图像生成完成！")
